import os
import shutil

# 设置路径
images_folder = r"E:\zhihao\2025文档\25年10月\ai_训练\ai项目_备份\yangzhihao\ai_project\public_security\traffic_accient\train\traffic_accidents_model\datasets\train\images"
labels_folder = r"E:\zhihao\2025文档\25年10月\ai_训练\ai项目_备份\yangzhihao\ai_project\public_security\traffic_accient\train\traffic_accidents_model\datasets\train\labels"

# images 文件夹中的 lyh 和 hx 文件夹路径
images_lyh_folder = os.path.join(images_folder, "lyh")
images_hx_folder = os.path.join(images_folder, "hx")

# 创建 labels 文件夹中的 lyh 和 hx 文件夹
labels_lyh_folder = os.path.join(labels_folder, "lyh")
labels_hx_folder = os.path.join(labels_folder, "hx")

os.makedirs(labels_lyh_folder, exist_ok=True)
os.makedirs(labels_hx_folder, exist_ok=True)

print("=" * 60)
print("开始同步 labels 文件夹的文件划分...")
print(f"✅ 已创建: {labels_lyh_folder}")
print(f"✅ 已创建: {labels_hx_folder}")
print("=" * 60)


def get_filename_without_ext(filename):
    """获取不带扩展名的文件名"""
    return os.path.splitext(filename)[0]


def copy_corresponding_labels(image_folder, label_source_folder, label_dest_folder, person_name):
    """根据图片文件名复制对应的标签文件"""

    # 获取图片文件夹中的所有文件
    image_files = [f for f in os.listdir(image_folder) if os.path.isfile(os.path.join(image_folder, f))]

    copied_count = 0
    not_found_count = 0
    not_found_files = []

    for img_file in image_files:
        # 获取不带扩展名的文件名
        base_name = get_filename_without_ext(img_file)

        # 在 labels 文件夹中查找对应的标签文件(可能是 .txt, .xml, .json 等)
        found = False
        for label_file in os.listdir(label_source_folder):
            label_path = os.path.join(label_source_folder, label_file)
            if os.path.isfile(label_path):
                label_base_name = get_filename_without_ext(label_file)
                if label_base_name == base_name:
                    # 找到对应的标签文件,复制它
                    dst = os.path.join(label_dest_folder, label_file)
                    shutil.copy2(label_path, dst)
                    copied_count += 1
                    found = True
                    break

        if not found:
            not_found_count += 1
            not_found_files.append(img_file)

    print(f"\n{person_name} 的统计:")
    print(f"  - 图片文件数: {len(image_files)}")
    print(f"  - 成功复制标签: {copied_count}")
    print(f"  - 未找到对应标签: {not_found_count}")

    if not_found_files and not_found_count <= 10:
        print(f"  - 未找到标签的图片:")
        for f in not_found_files[:10]:
            print(f"      {f}")

    return copied_count, not_found_count


# 处理 lyh 的文件
print("\n处理 lyh 的文件...")
lyh_copied, lyh_not_found = copy_corresponding_labels(
    images_lyh_folder,
    labels_folder,
    labels_lyh_folder,
    "lyh"
)

# 处理 hx 的文件
print("\n处理 hx 的文件...")
hx_copied, hx_not_found = copy_corresponding_labels(
    images_hx_folder,
    labels_folder,
    labels_hx_folder,
    "hx"
)

# 总结
print("\n" + "=" * 60)
print("✅ 同步完成!")
print(f"总结:")
print(f"  - lyh: 复制了 {lyh_copied} 个标签文件")
print(f"  - hx: 复制了 {hx_copied} 个标签文件")
print(f"  - 总共: {lyh_copied + hx_copied} 个标签文件")
if lyh_not_found + hx_not_found > 0:
    print(f"  ⚠️  警告: 有 {lyh_not_found + hx_not_found} 个图片文件没有找到对应的标签文件")
print("=" * 60)